#include <signal.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include "sanrazor.h"

static BOOL not_registered = true;
static struct SanSrcFileCovInfo* info_list_head = NULL;
static struct SanSrcFileCovInfo* info_list_tail = NULL;

void __asan_init();

char *get_dump_filename() {
  char *dump_path = getenv("SANRAZOR_DUMP");
  if (dump_path == NULL) {
    dump_path = (char *)malloc(1024);
    snprintf(dump_path, 1024, "%s", "sanrazor.log");
  }

  size_t name_size = strlen(dump_path);
  size_t log_id = 0;
  while (access(dump_path, F_OK) == 0) {
    snprintf(dump_path + name_size, 1024 - name_size, ".%zu", log_id);
    log_id++;
  }
  return dump_path;
}

void __init_san_cov_srcfile(struct SanSrcFileCovInfo* info) {
  if (info_list_head == NULL) {
    info_list_head = info;
    info_list_tail = info;
  } else {
    info_list_tail->next = info;
    info->next = NULL;
    info_list_tail = info;
  }
}

void __san_free_info_list() {
  struct SanSrcFileCovInfo* info = info_list_head;
  struct SanSrcFileCovInfo* next = NULL;
  while (info != NULL) {
    next = info->next;
    free(info);
    info = next;
  }
  info_list_head = NULL;
  info_list_tail = NULL;
}

char* __san_cov_internal_ulltoa(unsigned long long value, char* str) {
  static char digits[] = "0123456789abcdef";
  size_t idx = 0;
  while (value > 0) {
    unsigned long long tmp = (value & 0x0f);
    str[idx++] = digits[tmp];
    value >>= 4;
  }
  if (idx == 0) {
    str[0] = '0';
    str[1] = '\0';
    return str;
  }
  str[idx] = '\0';
  size_t len = 0;
  for (char c = str[len]; c != '\0'; c = str[++len]);
  // reverse str
  for (size_t i = 0; i < len / 2; ++i) {
    char tmp = str[i];
    str[i] = str[len - i - 1];
    str[len - i - 1] = tmp;
  }
  return str;
}

void __san_cov_dump() {
  struct SanSrcFileCovInfo* info = info_list_head;
  if (info == NULL) {
    return;
  }
  char *dump_path = get_dump_filename();
  fprintf(stderr, "[+] ASan report triggered - dumping log to %s\n", dump_path);
  FILE *fp = fopen(dump_path, "w");
  char buff[100];
  while (info != NULL) {
    for (size_t id = 0; id < info->size; ++id) {
      // Do not use snprintf as it is intercepted by ASan
      // snprintf(buff, 5000, "%s // %d // %d\n", info->filename, id, info->array[id]);
      // fwrite(buff, strlen(buff), 1, fp);
      fputs(info->filename, fp);
      fputs(" // ", fp);
      __san_cov_internal_ulltoa(id, buff);
      fputs(buff, fp);
      fputs(" // ", fp);
      __san_cov_internal_ulltoa(info->array[id], buff);
      fputs(buff, fp);
      fputs("\n", fp);
    }
    info = info->next;
  }
  fclose(fp);
  free(dump_path);
  __san_free_info_list();
}

void __san_cov_signal_handler(int sig) {
  signal(sig, SIG_DFL);
  __san_cov_dump();
  raise(sig);
}

void __san_cov_register() {
  if (!not_registered) {
    return;
  }
  /*Register the call to be called during function exit*/
  atexit(__san_cov_dump);
  signal(SIGINT, __san_cov_signal_handler);
  signal(SIGTERM, __san_cov_signal_handler);
  signal(SIGABRT, __san_cov_signal_handler);
  signal(SIGQUIT, __san_cov_signal_handler);
  signal(SIGSEGV, __san_cov_signal_handler);
  signal(SIGILL, __san_cov_signal_handler);
  signal(SIGFPE, __san_cov_signal_handler);
  signal(SIGBUS, __san_cov_signal_handler);
  not_registered = false;
  fprintf(stderr, "[+] ASan coverage report registered\n");
}

void __san_cov_abort() {
  fprintf(stderr, "[+] ASan report triggered - flushing log and aborting\n");
  __san_cov_dump();
  abort();
}

void __san_cov_append_info(uint64_t* array, BOOL *is_initialized, char* filename, size_t size) {
  // once the function is called before __asan_init
  __asan_init();
  struct SanSrcFileCovInfo* info = (struct SanSrcFileCovInfo*)malloc(sizeof(struct SanSrcFileCovInfo));
  info->filename = filename;
  info->array = array;
  info->size = size;
  info->next = NULL;
  __init_san_cov_srcfile(info);
  *is_initialized = true;
  // fprintf(stderr, "[+] ASan coverage registered - %s\n", filename);
}

void __san_cov_trace_pc(uint64_t* array, size_t id, BOOL *is_initialized, char* filename, size_t size) {
  if (!*is_initialized) {
    __san_cov_append_info(array, is_initialized, filename, size);
  }
  array[id]++;
}

void __san_cov_trace_pc2(uint64_t* array, size_t id) {
  array[id]++;
}